#!/usr/bin/python

### Compute all pairwise distances of patches.

import numpy,os,os.path,sys
from collections import Counter,defaultdict
from pylab import *
import tables,scipy
from scipy import fftpack
from scipy.cluster import vq
from scipy import stats
from scipy.ndimage import morphology,filters,interpolation
from ocrolib import docproc,improc
import pyflann
import random
from tables import openFile,Filters,Int32Atom,Float32Atom,Int64Atom
from multiprocessing import Pool
import fcntl,time
from llpy.sutils import *

import argparse
parser = argparse.ArgumentParser(description = "Compute character/line relation statistics for each prototype.")
parser.add_argument('inputs',nargs="*",help="input files")
parser.add_argument('-p','--protos',help="prototype file")
parser.add_argument('-o','--output',help="output file")
parser.add_argument('-N','--nsamples',default=2000000000,type=int,help="max # of samples")
args = parser.parse_args()

protos = Protos()
protos.load(args.protos)

rels = [[] for i in range(protos.n)]

for fname in args.inputs:
    with openFile(fname) as db:
        nsamples = min(args.nsamples,len(db.root.classes))
        for i in range(nsamples):
            if i%1000==0: print fname,i
            cls = db.root.classes[i]
            image = db.root.patches[i]
            rel = db.root.rels[i]
            [[r0]],[[d0]] = protos.nn_index(array([image.ravel()]))
            c0 = protos.classes[r0]
            if c0==cls: rels[r0].append(rel)

rmeans = []
rsigmas = []
for i in range(protos.n):
    if len(rels[i])==0:
        means = zeros(3)
        sigmas = zeros(3)
    else:
        means = mean(rels[i],axis=0)
        sigmas = sqrt(var(rels[i],axis=0))
    rmeans.append(means)
    rsigmas.append(sigmas)

with openFile(args.output,"w") as odb:
    with openFile(args.protos) as pdb:
        table_copy(pdb,odb)
        table_lcopy(pdb,odb)
    table_log(odb,"%s %s"%(sys.argv,time.asctime()))
    table_assign(odb,"relmeans",array(rmeans,'f'))
    table_assign(odb,"relsigmas",array(rsigmas,'f'))

